#coding:utf-8
import threading
import os
import traceback
import demjson
import requests
import time
import pycurl
from io import BytesIO
from OKHTTP import OKClass
ok = OKClass()
import globalvar
import datetime
#屏蔽requests不验证SSL证书警告报错
requests.packages.urllib3.disable_warnings(requests.packages.urllib3.exceptions.InsecureRequestWarning)
class spiderClass():
    '蜘蛛'
    def __init__(self, taskInfo):
        '初始化'
        self.Path = taskInfo['Path']
        self.IPFileUrl =  taskInfo['IPFileUrl']
        self.ScanType =  taskInfo['ScanType']
        self.ContainList =  taskInfo['ContainList']
        self.ExcludeList =  taskInfo['ExcludeList']
        self.TimeOut =  taskInfo['TimeOut']
        self.Thread =  taskInfo['Thread']
        self.HeadersExt =  taskInfo['Headers']
        self.SSL =  taskInfo['SSL']
        self.charset = taskInfo['charset']
        self.TASKID = taskInfo['TASKID']
        globalvar.set_value('TASKID', taskInfo['TASKID'])
        self.ResultUrl = taskInfo['ResultUrl']
        self.IPFilePath = 'IP.txt'
        self.Headers = {}
        self.taskInfo = {}
        self.AllowRedirects = taskInfo['AllowRedirects']
        self.Total = 0
        self.Now = 0
        self.UsedTime = 1
        self.AvgTime = 0
        self.NeedTime = 0
        self.SuccessCount = 0
        self.FailureCount = 0
        self.StartTime = 0
        self.EndTime = 0
        self.Found = 0
        self.SaveFoundPath = 'static/'+str(taskInfo['TASKID'])+'.txt'
        self.threads = []
        self.GroupAvg = 10000

        if os.path.exists('static') == False:
            os.mkdir('static')
    def SetHeaders(self):
        '设置Headers信息'
        try:
            tmpHeaders = demjson.decode(self.HeadersExt)
            for key in tmpHeaders.keys():
                self.Headers[key] = tmpHeaders[key]
            return True
        except Exception as e:
            print('SetHeaders:',e)
            return False
    def Filter(self,text):
        '过滤数据'
        try:
            result = True
            for Contain in self.ContainList:
                if Contain not in text:
                    result = False
            for Exclude in self.ExcludeList:
                if Exclude in text:
                    result = False
            del text
            return result
        except Exception as e:
            print('Filter:',e)
            return False
    def Save(self,path,text):
        saveFile = open(path,mode='a+')
        saveFile.write(text+'\n')
        saveFile.close()
    def GetIPFile(self):
        '下载IP文件'
        try:
            ipf = open(self.IPFilePath,'wb')
            ipf.write(requests.get(self.IPFileUrl).content)
            ipf.close()
            return True
        except Exception as e:
            print('GetIPFile:',e)
            return False
    def Grab(self,IP):
        '抓取程序'
        try:
            if self.SSL == True:
                url = 'https://' + IP + self.Path
            else:
                url = 'http://' + IP + self.Path
            if self.ScanType == 'GET':
                response = ok.GET(url=url,headers=self.Headers,AllowRedirects=self.AllowRedirects,timeout=self.TimeOut,charset=self.charset)
            else:
                response = ok.HEAD(url=url,headers=self.Headers,AllowRedirects=self.AllowRedirects,timeout=self.TimeOut,charset=self.charset)
            if response != False:
                if self.Filter(response) == True:
                    self.Save(self.SaveFoundPath,IP) #只保存IP
                    self.Found += 1
                    #上报至控制台
                    requests.post(self.ResultUrl,data={
                        'TaskID':self.TASKID,
                        'Result':IP
                    })
                self.SuccessCount += 1
            else:
                self.FailureCount += 1
            del response
        except Exception  as e:
            traceback.print_exc()
            self.FailureCount += 1
        self.UsedTime  = int(time.time() - self.StartTime) + 1
        self.Now += 1

        TaskStatusInfo = {
            'Total': self.Total,
            'Current': self.Now,
            'UsedTime': self.UsedTime,
            'Speed': str(int(self.Now / self.UsedTime)) + 'Req/S',
            'NeedTime': int((self.Total - self.Now) / int(self.Now / self.UsedTime)),
            'SuccessCount': self.SuccessCount,
            'FailureCount': self.FailureCount,
            'Found': self.Found,
            'UpdateTime': datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
            'status': globalvar.get_value('status')
        }

        globalvar.set_value('TaskStatusInfo', TaskStatusInfo)

    def Run(self):
        self.GetIPFile()
        self.SetHeaders()
        HostsFile = open(self.IPFilePath)
        Hosts = HostsFile.read().replace('\r','').split('\n')
        self.Total = len(Hosts)
        HostsFile.close()
        self.StartTime = time.time()
        GroupCount = int((self.Total - self.Total % self.GroupAvg) / self.GroupAvg)
        Remainder = self.Total % GroupCount
        self.GroupAvg = int((self.Total - Remainder) / GroupCount)
        for i in range(1, GroupCount + 1):
            GroupStart = (i - 1) * self.GroupAvg
            if i == GroupCount:
                tempHostList = Hosts[GroupStart:GroupStart + self.Total]
            else:
                tempHostList = Hosts[GroupStart:GroupStart + self.GroupAvg]
            for tempHost in tempHostList:
                t = threading.Thread(target=self.Grab, args=(tempHost,))
                self.threads.append(t)
            for thread in self.threads:
                thread.start()
                while True:
                    if len(threading.enumerate()) <= int(self.Thread):
                        break
            self.threads = []
        #锁住服务器状态
        globalvar.set_value('status', 'STOP')
